"""
基于车道中心线的曲线排队长度计算模块
"""

import numpy as np
import cv2
import logging
from typing import List, Dict, Any
from scipy.interpolate import CubicSpline
from scipy.integrate import quad

class CurvilinearLengthCalculator:
    """曲线长度计算器"""
    
    def __init__(self, config):
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 车道几何信息
        self.centerline_points = np.array(config.centerline_points or [])
        self.stop_line_position = np.array(config.stop_line_position or [0, 0])
        self.camera_matrix = np.array(config.camera_matrix) if config.camera_matrix else None
        self.dist_coeffs = np.array(config.dist_coeffs) if config.dist_coeffs else None
        
        # 样条曲线拟合
        self.spline = self.fit_spline_curve()
        
    def fit_spline_curve(self):
        """拟合样条曲线"""
        if len(self.centerline_points) < 4:
            self.logger.warning("中心线点数量不足，使用线性插值")
            return None
            
        try:
            # 参数化曲线
            points = np.array(self.centerline_points)
            cumulative_dist = self.calculate_cumulative_distance(points)
            t_normalized = cumulative_dist / cumulative_dist[-1]
            
            # 三次样条拟合
            cs_x = CubicSpline(t_normalized, points[:, 0])
            cs_y = CubicSpline(t_normalized, points[:, 1])
            
            return {'x': cs_x, 'y': cs_y, 't_range': [0, 1]}
            
        except Exception as e:
            self.logger.error(f"样条拟合失败: {str(e)}")
            return None
    
    def calculate_cumulative_distance(self, points):
        """计算累积距离"""
        if len(points) < 2:
            return np.array([0])
            
        distances = [0]
        for i in range(1, len(points)):
            dist = np.linalg.norm(points[i] - points[i-1])
            distances.append(distances[-1] + dist)
            
        return np.array(distances)
    
    def calculate_length(self, cluster):
        """计算排队长度"""
        if not cluster:
            return 0.0
            
        # 找到队尾车辆
        tail_vehicle = self.find_tail_vehicle(cluster)
        if not tail_vehicle:
            return 0.0
            
        # 获取队尾车辆位置
        tail_position = self.get_vehicle_position(tail_vehicle)
        if tail_position is None:
            return 0.0
            
        # 计算曲线距离
        if self.spline is not None:
            curve_length = self.calculate_curvilinear_distance(
                tail_position, self.stop_line_position)
        else:
            # 备用：直线距离
            curve_length = np.linalg.norm(tail_position - self.stop_line_position)
        
        # 坐标转换（像素到米）
        physical_length = curve_length * self.config.pixel_to_meter_ratio
        
        return physical_length
    
    def find_tail_vehicle(self, cluster):
        """找到队尾车辆"""
        if not cluster:
            return None
            
        # 计算每辆车到停车线的距离
        max_distance = -1
        tail_vehicle = None
        
        for vehicle in cluster:
            position = self.get_vehicle_position(vehicle)
            if position is None:
                continue
                
            if self.spline is not None:
                distance = self.calculate_curvilinear_distance(
                    position, self.stop_line_position)
            else:
                distance = np.linalg.norm(position - self.stop_line_position)
                
            if distance > max_distance:
                max_distance = distance
                tail_vehicle = vehicle
                
        return tail_vehicle
    
    def get_vehicle_position(self, vehicle):
        """获取车辆位置"""
        # 从车辆轨迹中获取最新位置
        trajectory = vehicle.get('trajectory', [])
        if trajectory:
            return np.array(trajectory[-1])
        
        # 从检测框中获取位置
        if 'last_detection' in vehicle:
            bbox = vehicle['last_detection'].get('bbox', [])
            if len(bbox) >= 2:
                return np.array([bbox[0], bbox[1]])
                
        return None
    
    def calculate_curvilinear_distance(self, point_a, point_b):
        """计算两点间的曲线距离"""
        if self.spline is None:
            return np.linalg.norm(point_a - point_b)
            
        try:
            # 找到点在曲线上的最近参数
            t_a = self.find_nearest_parameter(point_a)
            t_b = self.find_nearest_parameter(point_b)
            
            # 确保参数顺序
            t_start, t_end = sorted([t_a, t_b])
            
            # 数值积分计算弧长
            length, error = quad(self.arc_length_integrand, t_start, t_end)
            return length
            
        except Exception as e:
            self.logger.warning(f"曲线距离计算失败: {str(e)}")
            return np.linalg.norm(point_a - point_b)
    
    def find_nearest_parameter(self, point):
        """找到点在曲线上的最近参数"""
        if self.spline is None:
            return 0.0
            
        from scipy.optimize import minimize_scalar
        
        def distance_to_curve(t):
            curve_point = np.array([
                self.spline['x'](t),
                self.spline['y'](t)
            ])
            return np.linalg.norm(point - curve_point)
        
        try:
            result = minimize_scalar(
                distance_to_curve, 
                bounds=self.spline['t_range'], 
                method='bounded'
            )
            return result.x
        except:
            return 0.5  # 默认参数
    
    def arc_length_integrand(self, t):
        """弧长被积函数"""
        if self.spline is None:
            return 1.0
            
        try:
            dx_dt = self.spline['x'](t, 1)  # 一阶导数
            dy_dt = self.spline['y'](t, 1)
            return np.sqrt(dx_dt**2 + dy_dt**2)
        except:
            return 1.0
    
    def undistort_points(self, points):
        """去除图像畸变"""
        if self.camera_matrix is None or self.dist_coeffs is None:
            return points
            
        try:
            points = np.array(points, dtype=np.float32).reshape(-1, 1, 2)
            undistorted = cv2.undistortPoints(
                points, self.camera_matrix, self.dist_coeffs, P=self.camera_matrix)
            return undistorted.reshape(-1, 2)
        except Exception as e:
            self.logger.warning(f"去畸变失败: {str(e)}")
            return points
    
    def calculate_multiple_cluster_lengths(self, clusters):
        """计算多个集群的排队长度"""
        lengths = {}
        
        for i, cluster in enumerate(clusters):
            length = self.calculate_length(cluster)
            lengths[f'cluster_{i}'] = {
                'length': length,
                'vehicle_count': len(cluster),
                'tail_vehicle_id': self.get_tail_vehicle_id(cluster)
            }
        
        return lengths
    
    def get_tail_vehicle_id(self, cluster):
        """获取队尾车辆ID"""
        tail_vehicle = self.find_tail_vehicle(cluster)
        return tail_vehicle.get('track_id', -1) if tail_vehicle else -1
    
    def validate_geometry(self):
        """验证几何配置"""
        issues = []
        
        if len(self.centerline_points) < 2:
            issues.append("中心线点数量不足")
            
        if len(self.stop_line_position) != 2:
            issues.append("停车线位置格式错误")
            
        if self.config.pixel_to_meter_ratio <= 0:
            issues.append("像素到米转换比例必须为正数")
            
        return issues